# -*- coding: utf-8 -*-
# -*- coding:utf-8 -*-
import csv
import os

from options import HiDDenConfiguration
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torchvision
import utils
from model.discriminator import Discriminator
import numpy as np
from model.encoder_decoder import EncoderDecoder
from modules.UnetPlusPlus_L6 import UNetPlusPlusL6_Attack
from noise_layers.jpeg import Jpeg
from noise_layers.colorjitter import ColorJitter
from model.attack import Attack
import config as c
from noise_layers.gaussian_noise import Gaussian_Noise
from noise_layers.salt_pepper_noise import SP
from modules.RandompixelUNet import DistortionNet
from noise_layers.poisson import PoissonNoise

def model_save(model, ep, save_dir):
    save_path = c.model_dir
    print(f'--- save the model @ epoch {ep} ---')

    checkpoint = {
        'epoch': ep,
        'encoder_decoder_state_dict': model.encoder_decoder.state_dict(),
        'dis_state_dict': model.discriminator.state_dict(),
        # 'noiser_state_dict': model.noise_layer.state_dict(),
        # 'ed_opt_state_dict': model.optimizer_ed.state_dict(),
        # 'dis_opt_state_dict': model.optimizer_dis.state_dict(),
        # 'noiser_opt_state_dict': model.optimizer_n.state_dict(),
        # 'log_sigma_asl': model.log_sigma_asl.detach().cpu(),
        # 'log_sigma_sml': model.log_sigma_sml.detach().cpu(),
    }

    torch.save(checkpoint, save_dir)


class  ED_NR(nn.Module):
    def __init__(self):
        super(ED_NR, self).__init__()
        hidden_config = HiDDenConfiguration(H=128, W=128,
                                        message_length=c.message_length,
                                        encoder_blocks=4, encoder_channels=64,
                                        decoder_blocks=7, decoder_channels=64,
                                        use_discriminator=True,
                                        use_vgg=False,
                                        discriminator_blocks=3, discriminator_channels=64,
                                        decoder_loss=1,
                                        encoder_loss=0.7,
                                        adversarial_loss=1e-3
                                        )
        #################
        ###   Model   ###
        #################
        self.discriminator = Discriminator(hidden_config).cuda()
        self.encoder_decoder = EncoderDecoder(hidden_config).cuda()
        # self.noise_layer = ParallelImageNet(batch_size=c.batch_size, image_size=c.image_size, num_layers=128).cuda()
        # self.noise_layer = UNetPlusPlusL7_Attack().cuda()
        self.noise_layer = DistortionNet().cuda()
        # self.noise_layer = UNet_Attack().cuda()
        self.log_sigma_asl = nn.Parameter(torch.zeros(1, device="cuda"))
        self.log_sigma_sml = nn.Parameter(torch.zeros(1, device="cuda"))

        #################
        ###    opt    ###
        #################
        log10_lr = -4.0
        lr = 10 ** log10_lr
        self.optimizer_ed = optim.Adam(
            list(self.encoder_decoder.parameters()) + [self.log_sigma_asl, self.log_sigma_sml],
            lr=lr
        )
        self.optimizer_dis = optim.Adam(self.discriminator.parameters(), lr=0.0001)
        self.optimizer_n = optim.Adam(self.noise_layer.parameters(), lr=0.0001)

        ################
        ###   Loss   ###
        ################
        
        self.mse_loss = nn.MSELoss().cuda()
        self.L1 = nn.L1Loss().cuda()
        self.bce_with_logits_loss = nn.BCEWithLogitsLoss()

        ######################
        ###   Pretrained   ###
        ######################
        self.trained = False
        self.epoch = 0
        if not self.trained:
            self.encoder_decoder.load_state_dict(torch.load("/home/zzc/simulator_noisy/run_base_2025-04-10_14-11-47/model/300_model.pth"))
        else:
            state_dict = torch.load("/home/zzc/simulator_noisy/HiDDeN-ED-Random/runs-works/run_Brightness-no-normalize_sampled_100_en0.7_de1.0_2025-07-22_16-07-58/model/best.pth")
            self.epoch = state_dict['epoch']
            self.encoder_decoder.load_state_dict(state_dict['encoder_decoder_state_dict'])
            self.discriminator.load_state_dict(state_dict['dis_state_dict'])
            self.noise_layer.load_state_dict(state_dict['noiser_state_dict'])
            self.optimizer_ed.load_state_dict(state_dict['ed_opt_state_dict'])
            for param_group in self.optimizer_ed.param_groups:
                param_group['lr'] = 0.00001 # 降低学习率
            self.optimizer_dis.load_state_dict(state_dict['dis_opt_state_dict'])
            for param_group in self.optimizer_dis.param_groups:
                param_group['lr'] = 0.00001 # 降低学习率
            self.optimizer_n.load_state_dict(state_dict['noiser_opt_state_dict'])
            for param_group in self.optimizer_n.param_groups:
                param_group['lr'] = 0.00001 # 降低学习率
        print("This is start as {} epooch".format(self.epoch))
        #################
        ###   noise   ###
        #################
        # distortion = 'Brightness'   # 亮度
        # distortion = 'Contrast'   # 对比度
        # distortion = 'Hue'   # 色相
        # distortion = 'Saturation'   # 饱和度
        # self.ASL = ColorJitter(distortion).cuda()
        # self.ASL = Gaussian_Noise().cuda()
        self.ASL = SP().cuda()
        # self.ASL = Jpeg(Q=90).cuda()
        # self.ASL = PoissonNoise().cuda()
            


        self.train_noiser_loss = 0.0
        self.train_ED_loss_ASL = 0.0
        self.train_ED_loss_SML = 0.0
        self.train_ED_loss = 0.0

        with open(c.log_file, "w") as f:
            writer = csv.writer(f)
            writer.writerow(["Epoch", "trainLoss", "train_loss_ASL", "train_loss_SML", 
                    "testLoss", "test_encoder_loss", "test_decoder_loss", 
                    "advLoss", "bit_arr", "psnr", "ssim", "train_noiser_loss", "test_noiser_loss",
                    "w_asl", "w_sml"])
        self.min_testLoss = float("INF")

    def res_init(self):

        self.train_noiser_loss = 0.0
        self.train_ED_loss_ASL = 0.0
        self.train_ED_loss_SML = 0.0
        self.train_ED_loss = 0.0
        
    def update_encoder_decoder_ASL_And_SML(self, images):

        cover_data_label = 1
        encoder_data_label = 0
        self.encoder_decoder.train()
        self.discriminator.train()
        self.noise_layer.eval()


        images = images.cuda()
            

        target_label_cover = torch.full((images.size(0), 1), cover_data_label).cuda()
        target_label_encoded = torch.full((images.size(0), 1), encoder_data_label).cuda()
        target_label_adv = torch.full((images.size(0), 1), cover_data_label).cuda()
        messages = torch.Tensor(np.random.choice([0, 1], size=[images.size(0), 1, c.message_length, c.message_length])).cuda()
        
        # 修改Forward ASL部分，添加数值保护
        with torch.enable_grad():
            encoder_image = self.encoder_decoder.encoder(images, messages)
            encoder_image = torch.clamp(encoder_image, 0, 1)  # 确保编码输出在[0,1]范围内

        with torch.no_grad():
            attack_image = self.ASL(encoder_image)
            attack_image = torch.clamp(attack_image, 0, 1)  # 确保ASL输出在有效范围内
            gap_image = attack_image - encoder_image
            gap_image = gap_image.detach()

        with torch.enable_grad():
            noise_image = encoder_image + gap_image
            noise_image = torch.clamp(noise_image, 0, 1)  # 确保噪声图像在有效范围内
            decoder_message = self.encoder_decoder.decoder(noise_image)
            decoder_message = torch.clamp(decoder_message, 0, 1)  # 确保解码输出在[0,1]范围内
    
        with torch.enable_grad():

            # 训练discriminator
            self.optimizer_dis.zero_grad()
            adv_loss_on_cover = self.bce_with_logits_loss(self.discriminator(images), target_label_cover.float())
            adv_loss_on_cover.backward()

            adv_loss_on_encoder = self.bce_with_logits_loss(self.discriminator(encoder_image.detach()), target_label_encoded.float())
            adv_loss_on_encoder.backward()
            self.optimizer_dis.step()

            # 训练encoder_decoder
            self.optimizer_ed.zero_grad()
            adv_loss_ASL = self.bce_with_logits_loss(self.discriminator(encoder_image), target_label_adv.float()) * 100
            encoder_loss_ASL = self.mse_loss(encoder_image, images) * 100
            decoder_loss_ASL = self.mse_loss(decoder_message, messages) * 100

            lossASL = c.en * encoder_loss_ASL + c.de * decoder_loss_ASL + c.adv * adv_loss_ASL

            encoder_image = self.encoder_decoder.encoder(images.clone(), messages.clone())
            encoder_image = torch.clamp(encoder_image, 0, 1)  # 确保编码输出在[0,1]范围内
            noise_image = self.noise_layer(encoder_image)
            noise_image = torch.clamp(noise_image, 0, 1)
            decoder_message = self.encoder_decoder.decoder(noise_image)
            decoder_message = torch.clamp(decoder_message, 0, 1)

            adv_loss_SML = self.bce_with_logits_loss(self.discriminator(encoder_image), target_label_adv.float()) * 100
            encoder_loss_SML = self.mse_loss(encoder_image, images) * 100
            decoder_loss_SML = self.mse_loss(decoder_message, messages) * 100

            lossSML = c.en * encoder_loss_SML + c.de * decoder_loss_SML + c.adv * adv_loss_SML

            loss = (torch.exp(-2 * self.log_sigma_asl)  * lossASL +
                    torch.exp(-2 * self.log_sigma_sml) * lossSML +
                    (self.log_sigma_asl + self.log_sigma_sml))

            loss.backward()
            self.optimizer_ed.step()

            self.train_ED_loss_ASL += lossASL.item()
            self.train_ED_loss_SML += lossSML.item()
            self.train_ED_loss += loss.item()


    def update_noiser(self, images):
        self.noise_layer.train()
        self.encoder_decoder.eval()
        messages = torch.Tensor(np.random.choice([0, 1], size=[images.size(0), 1, c.message_length, c.message_length])).cuda()
        stego = self.encoder_decoder.encoder(images, messages)
        stego = torch.clamp(stego, 0, 1)
        with torch.enable_grad():
            noise_images = self.ASL(stego)
            self.optimizer_n.zero_grad()
            attack_images = self.noise_layer(stego)

            loss = self.mse_loss(attack_images, noise_images)
            loss.backward()
            self.optimizer_n.step()
            self.train_noiser_loss += loss.item()

    def test_EDN(self, ep_num, iter_num, test_loader, isSave = True):
        train_epoch_loss_ASL = self.train_ED_loss_ASL / iter_num
        train_epoch_loss_SML = self.train_ED_loss_SML / iter_num
        train_epoch_loss = self.train_ED_loss / iter_num
        train_epoch_noiser_loss = self.train_noiser_loss / iter_num

        is_first_save = True
        test_noiser_loss = 0.0
        test_encoder_loss = 0.0
        test_decoder_loss = 0.0
        test_psnr = 0.0
        test_ssim = 0.0
        test_bit_arr = 0.0
        test_running_loss = 0.0
        dis_running_loss = 0.0
        cover_data_label = 1
        # 测试阶段
        self.encoder_decoder.eval()
        self.noise_layer.eval()
        with torch.no_grad():
            with open(c.log_file, "a") as f:
                writer = csv.writer(f)
                for images, _ in test_loader:
                    images = images.cuda()
                    target_label_adv = torch.full((images.size(0), 1), cover_data_label).cuda()

                    messages = torch.Tensor(np.random.choice([0, 1], size=[images.size(0), 1, c.message_length, c.message_length])).cuda()
                    encoder_image = self.encoder_decoder.encoder(images, messages)
                    encoder_image = torch.clamp(encoder_image, 0, 1)
                    noise_image = self.ASL(encoder_image)
                    noise_image = torch.clamp(noise_image, 0, 1)
                    attack_image = self.noise_layer(encoder_image)
                    attack_image = torch.clamp(attack_image, 0, 1)
                    decoder_message = self.encoder_decoder.decoder(noise_image)
                    
                    noiser_loss = self.mse_loss(attack_image, noise_image) 
                    dis_loss = self.bce_with_logits_loss(self.discriminator(encoder_image), target_label_adv.float())  * 100
                    encoder_loss = self.mse_loss(images, encoder_image)  * 100
                    decoder_loss = self.mse_loss(messages, decoder_message) * 100
                    loss = c.en * encoder_loss + c.de * decoder_loss + c.adv * dis_loss

                    psnr, ssim = utils.val_psnr_ssim(images, encoder_image)
                    test_psnr += psnr
                    test_ssim += ssim

                    test_encoder_loss += encoder_loss.item()
                    test_decoder_loss += decoder_loss.item()
                    test_running_loss += loss.item()
                    test_noiser_loss += noiser_loss.item()
                    dis_running_loss += dis_loss.item()


                    if is_first_save:  # 只保存第一个批次的图像
                        # is_first_save = False
                        save_images(images.cpu()[:8, :, :, :],
                                    encoder_image[:8, :, :, :].cpu(), noise_image[:8, :, :, :].cpu(), ep_num, c.save_image_dir)
                    decoder_message = (decoder_message > 0.5).float()  # torch 中做稳定的二值化
                    bitwise_err = torch.abs((decoder_message > 0.5).float() - messages).mean().item()
                    test_bit_arr += bitwise_err

                test_epoch_encoder_loss = test_encoder_loss / len(test_loader)
                test_epoch_decoder_loss = test_decoder_loss / len(test_loader)
                test_epoch_loss = test_running_loss / len(test_loader)
                test_epoch_noiser_loss = test_noiser_loss / len(test_loader)
                dis_epoch_loss = dis_running_loss / len(test_loader)
                psnr = test_psnr / len(test_loader)
                ssim = test_ssim / len(test_loader)
                bit_arr = test_bit_arr / len(test_loader)

                writer.writerow([
                    ep_num, 
                    f"{train_epoch_loss:.4f}", 
                    f"{train_epoch_loss_ASL:.4f}", 
                    f"{train_epoch_loss_SML:.4f}", 
                    f"{test_epoch_loss:.4f}", 
                    f"{test_epoch_encoder_loss:.4f}", 
                    f"{test_epoch_decoder_loss:.4f}", 
                    f"{dis_epoch_loss:.4f}", 
                    f"{bit_arr:.4f}", 
                    f"{psnr:.2f}", 
                    f"{ssim:.2f}",
                    f"{train_epoch_noiser_loss:.4f}",
                    f"{test_epoch_noiser_loss:.4f}",
                    f"{torch.exp(-2*self.log_sigma_asl).item():.4f}",
                    f"{torch.exp(-2*self.log_sigma_sml).item():.4f}"
                
                ])
                f.flush()
                print(f"Epoch [{ep_num}/{c.epochs}], "
                                f"Train Loss: {train_epoch_loss:.4f} (Train ASL Loss: {train_epoch_loss_ASL:.4f}, Train SML Loss: {train_epoch_loss_SML:.4f}), "
                                f"Test Loss: {test_epoch_loss:.4f} (Encoder: {test_epoch_encoder_loss:.4f}, Decoder: {test_epoch_decoder_loss:.4f}), "
                                f"Adv Loss: {dis_epoch_loss:.4f}, Bit arr: {bit_arr:.4f}, Psnr: {psnr}, Ssim: {ssim}, "
                                f"Train noiser Loss: {train_epoch_noiser_loss:.4f}, Test noiser Loss:{test_epoch_noiser_loss:.4f}")
                print(f"w_asl={torch.exp(-2*self.log_sigma_asl).item():.4f}, "
                      f"w_sml={torch.exp(-2*self.log_sigma_sml).item():.4f}")

                if self.min_testLoss > test_epoch_loss:
                    self.min_testLoss = test_epoch_loss
                    model_save(self, ep_num, f'{c.model_dir}/best.pth')
                
                if isSave:
                    model_save(self, ep_num, f'{c.model_dir}/{ep_num}.pth')



def save_images(cover_images, encoder_images, noise_image, epoch, folder):
    # 将数据转移到 CPU
    images = cover_images.cpu()
    encoder_images = encoder_images.cpu()
    noise_image = noise_image.cpu()


    # 计算误差放大
    revise_image = torch.clamp(torch.abs(images - encoder_images) * 5, 0, 1)
    revise_image_5 = torch.clamp(torch.abs(images - noise_image) * 5, 0, 1)

    # 拼接图像
    stacked_images = torch.cat([images, encoder_images, noise_image, revise_image, revise_image_5], dim=0)

    # 保存图片
    filename = os.path.join(folder, 'epoch-{}.png'.format(epoch))
    torchvision.utils.save_image(stacked_images, filename, normalize=False)


